package com.zixi.ai.ai.service.impl;

import cn.hutool.core.bean.BeanUtil;
import com.alibaba.fastjson2.JSONObject;
import com.zixi.ai.ai.service.IAiModelService;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.OllamaEmbeddingModel;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.ollama.management.ModelManagementOptions;
import org.springframework.ai.ollama.management.PullModelStrategy;
import org.springframework.stereotype.Component;

import java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author wfg
 */
@RequiredArgsConstructor
@Component("ollamaAiModel")
public class OllamaAiModelService implements IAiModelService {
    private static final Map<String, EmbeddingModel> EMBEDDING_MODEL_MAP = new ConcurrentHashMap<>(16);
    private static final Map<String, ChatModel> CHAT_MODEL_MAP = new ConcurrentHashMap<>(16);

    @Override
    public EmbeddingModel buildEmbeddingModel(JSONObject config) {
        String baseUrl = config.getString("baseUrl");
        String model = config.getString("model");
        EmbeddingModel embeddingModel = EMBEDDING_MODEL_MAP.get(baseUrl);
        if (embeddingModel == null) {
            PullModelStrategy embeddingModelPullStrategy = PullModelStrategy.NEVER;
            OllamaApi ollamaApi = OllamaApi.builder().baseUrl(baseUrl).build();
            OllamaOptions options = new OllamaOptions();
            options.setModel(model);
            BeanUtil.fillBeanWithMap(config, options, true);
            embeddingModel = OllamaEmbeddingModel.builder()
                    .ollamaApi(ollamaApi)
                    .defaultOptions(options)
                    .modelManagementOptions(new ModelManagementOptions(embeddingModelPullStrategy,
                            List.of(), Duration.ofMinutes(5), 10))
                    .build();
            EMBEDDING_MODEL_MAP.put(baseUrl + ":" + model, embeddingModel);
        }
        return embeddingModel;
    }

    @Override
    public ChatModel buildChatModel(JSONObject config) {
        String baseUrl = config.getString("baseUrl");
        String model = config.getString("model");
        ChatModel chatModel = CHAT_MODEL_MAP.get(baseUrl);
        if (chatModel == null) {
            OllamaApi ollamaApi = OllamaApi.builder().baseUrl(baseUrl).build();
            OllamaOptions options = new OllamaOptions();
            options.setModel(model);
            BeanUtil.fillBeanWithMap(config, options, true);
            chatModel = OllamaChatModel.builder()
                    .ollamaApi(ollamaApi)
                    .defaultOptions(options)
                    .build();
            CHAT_MODEL_MAP.put(baseUrl + ":" + model, chatModel);
        }
        return chatModel;
    }
}
